#!/usr/bin/env Rscript
# -*- coding: utf-8 -*-
# ------------------------------------------------------------------------------
# @File    : alpha_diversity.R
# @Author  : Bing Liang
# @Date    : 2025/10/12
# @Desc    : 从 RPM 矩阵计算 Alpha 多样性，结合分组信息绘制柱状图和箱线图并做差异检验
# ------------------------------------------------------------------------------

suppressPackageStartupMessages({
  library(optparse)
  library(tidyverse)
  library(vegan)
  library(ggplot2)
  library(ggpubr)   # 用于显著性标记
  library(rstatix)  # 方便做多组差异检验
})

# ------------------------------------------------------------------------------
# 函数定义
# ------------------------------------------------------------------------------

mkdirs <- function(path) {
  dir.create(path, recursive = TRUE, showWarnings = FALSE)
}

# 读取 RPM 数据并转置
read_rpm_matrix <- function(file) {
  df <- read.table(file, header = TRUE, sep = "\t", row.names = 1, check.names = FALSE)
  t(df)  # 转置为样本为行，物种为列
}

# 读取分组配置文件（两列：Sample\tGroup）
read_group_file <- function(file) {
  read.table(file, header = TRUE, sep = "\t", stringsAsFactors = FALSE)
}

# 计算 alpha 多样性指标
calc_alpha_diversity <- function(data_t) {
  data.frame(
    Sample   = rownames(data_t),
    Shannon  = diversity(data_t, index = "shannon"),
    Simpson  = diversity(data_t, index = "simpson"),
    Richness = specnumber(data_t)
  )
}

save_alpha_table <- function(alpha_res, output_file) {
  mkdirs(dirname(output_file))
  write.table(alpha_res, output_file, sep = "\t", quote = FALSE, row.names = FALSE)
}



# 柱状图（按分组分面）
plot_alpha_bar <- function(alpha_res, group_df, outpdf, outpng, width = 20, height = 6, dpi = 360) {
  df <- alpha_res %>% left_join(group_df, by = "Sample")
  
  p <- df %>%
    pivot_longer(cols = c(Shannon, Simpson, Richness), names_to = "Index", values_to = "Value") %>%
    ggplot(aes(x = Sample, y = log2(Value + 1), fill = Index)) +
    geom_bar(stat = "identity", position = "dodge") +
    facet_wrap(~Group, scales = "free_x") +   # 按分组分面
    theme_bw() +
    labs(title = "Alpha Diversity (Barplot by Group)", y = "log2(value + 1)", x = "Sample") +
    theme(axis.text.x = element_text(angle = 270, vjust = 0.5))
  
  ggsave(plot = p, filename = outpdf, width = width, height = height)
  ggsave(plot = p, filename = outpng, width = width, height = height, dpi = dpi)
}

plot_alpha_box <- function(alpha_res, group_df, outpdf, outpng, width = 8, height = 6, dpi = 360) {
  df <- alpha_res %>% left_join(group_df, by = "Sample")
  
  # 转为长格式
  df_long <- df %>%
    pivot_longer(cols = c(Shannon, Simpson, Richness),
                 names_to = "Index", values_to = "Value")
  
  # 获取所有指数名称
  indices <- unique(df_long$Index)
  
  for (idx in indices) {
    df_idx <- df_long %>% filter(Index == idx)
    
    # --------------------------
    # 两两比较（BH校正）
    # --------------------------
    pwc <- df_idx %>%
      rstatix::pairwise_wilcox_test(Value ~ Group, p.adjust.method = "BH") %>%
      rstatix::add_significance("p") %>%
      rstatix::add_xy_position(x = "Group")
    
    # --------------------------
    # 自动调整显著性标记高度（避免重叠）
    # --------------------------
    if (nrow(pwc) > 0) {
      # 计算每个显著性标记的垂直间距
      y_range <- max(df_idx$Value, na.rm = TRUE) - min(df_idx$Value, na.rm = TRUE)
      step <- y_range * 0.1
      
      # 按照比较顺序依次上移显著性线条，避免重叠
      pwc$y.position <- pwc$y.position + seq(0, by = step, length.out = nrow(pwc))
    }
    
    # 绘图
    p <- ggboxplot(df_idx,
                   x = "Group", y = "Value",
                   fill = "Group",
                   palette = "Set2") +
      stat_pvalue_manual(pwc, label = "p.signif", tip.length = 0.01) +
      labs(title = paste("Alpha Diversity -", idx),
           y = idx, x = "Group") +
      theme_bw() +
      theme(
        plot.title = element_text(size = 14, face = "bold", hjust = 0.5),
        axis.text.x = element_text(angle = 45, hjust = 1)
      )
    
    # 输出文件（自动带上指标名）
    pdf_file <- gsub("\\.pdf$", paste0("_", idx, ".pdf"), outpdf)
    png_file <- gsub("\\.png$", paste0("_", idx, ".png"), outpng)
    
    ggsave(plot = p, filename = pdf_file, width = width, height = height)
    ggsave(plot = p, filename = png_file, width = width, height = height, dpi = dpi)
    
    message(paste("✅ 已保存指数", idx, "的箱线图：", pdf_file))
  }
}






# ------------------------------------------------------------------------------
# 主函数
# ------------------------------------------------------------------------------

main <- function() {
  option_list <- list(
    make_option(c("-i", "--input"), type = "character", help = "RPM矩阵（行为物种，列为样本）"),
    make_option(c("-m", "--meta"), type = "character", help = "分组配置文件（两列：Sample, Group）"),
    make_option(c("-o", "--output"), type = "character", help = "输出文件（指数矩阵，TSV格式）"),
    make_option(c("-f", "--outpdf"), type = "character", help = "柱状图PDF"),
    make_option(c("-g", "--outpng"), type = "character", help = "柱状图PNG"),
    make_option(c("-F", "--boxpdf"), type = "character", help = "箱线图PDF"),
    make_option(c("-G", "--boxpng"), type = "character", help = "箱线图PNG"),
    make_option(c("-W", "--width"), type = "numeric", default = 20.0, help = "柱状图宽度"),
    make_option(c("-H", "--height"), type = "numeric", default = 6.0, help = "柱状图高度")
  )
  
  opt <- parse_args(OptionParser(option_list = option_list))
  
  if (is.null(opt$input) || is.null(opt$meta) || is.null(opt$output) || 
      is.null(opt$outpdf) || is.null(opt$outpng) || is.null(opt$boxpdf) || is.null(opt$boxpng)) {
    stop("参数缺失，请确保 --input, --meta, --output, --outpdf, --outpng, --boxpdf, --boxpng 都已提供")
  }
  
  data_t <- read_rpm_matrix(opt$input)
  group_df <- read_group_file(opt$meta)
  alpha_res <- calc_alpha_diversity(data_t)
  
  save_alpha_table(alpha_res, opt$output)
  plot_alpha_bar(alpha_res, group_df, opt$outpdf, opt$outpng, opt$width, opt$height)
  plot_alpha_box(alpha_res, group_df, opt$boxpdf, opt$boxpng)
  
  message("Alpha 多样性计算、分组绘图与差异检验完成")
}

main()
